%load_ext autoreload
%autoreload 2

import banner

topics = ['Introduction to Transformers',
          'Very Simple Implementation of Transformers']
          

banner.reset(topics)

banner.next_topic()

banner.next_topic()

df_fake = pandas.read_csv('Fake.csv', usecols=['title'])
df_real = pandas.read_csv('True.csv', usecols=['title'])
df_fake.shape, df_real.shape

df_fake.head(5)

df_real.head(5)

keep = 1000
headlines_fake = df_fake.values[:keep]
labels_fake = numpy.zeros((headlines_fake.shape[0]))

headlines_real = df_real.values[:keep]
labels_real = numpy.ones((headlines_real.shape[0]))

headlines_fake.shape, labels_fake.shape, headlines_real.shape, labels_real.shape

from string import punctuation
punctuation + ' '
print(punctuation)

def clean_up_words(titles):
    import re
    titles_words = [[word.lower() for word in re.split('\W+', title[0])] for title in titles]
    # words = [w.strip(punctuation) for w in words]
    words = [[w for w in title if len(w) > 1] for title in titles_words] 
    return words

re.split('\W+', headlines_fake[0][0])

clean_up_words(headlines_fake[0:2])

headlines = headlines_fake + headlines_real
labels = numpy.hstack((labels_fake, labels_real))
len([len(h) for h in headlines]), len(labels)

mx = max([len(h) for h in headlines])
mx

len(headlines[0]), headlines[0]

words = [word for headline in headlines for word in headline]
vocabulary = numpy.unique(words)
len(vocabulary)

for i in range(100):
    print(i, vocabulary[i], end='; ')
print()
for i in range(4300, 4581):
    print(i, vocabulary[i], end='; ')

numpy.where(numpy.array(headlines[0]).reshape(-1, 1) == vocabulary)

numpy.where(numpy.array(headlines[0]).reshape(-1, 1) == vocabulary)[1]

tokens = [numpy.where(numpy.array(headline).reshape(-1, 1) == vocabulary)[1] for headline in headlines]
tokens[:20]

tokens[0]

X_tokens = jnp.array(tokens)
X_tokens

######################################################################
# 1. Make embedding layer

n_vocabulary_words = len(vocabulary)
embed_dim = 40  
embedder_W = jnp.array(numpy.random.normal(size=(n_vocabulary_words, embed_dim)))

type(embedder_W)  # , embedder_W.device()

######################################################################
# 2. Make positional encoding (from http://jalammar.github.io/illustrated-transformer/)

def get_angles(pos, i, d_model):
   angle_rates = 1 / numpy.power(10, (2 * (i//2)) / numpy.float32(d_model))
   return pos * angle_rates

def make_position_encoding(position, d_model):
   angle_rads = get_angles(numpy.arange(position)[:, numpy.newaxis],
                           numpy.arange(d_model)[numpy.newaxis, :],
                           d_model)
   # apply sin to even indices in the array; 2i
   angle_rads[:, 0::2] = numpy.sin(angle_rads[:, 0::2])
   # apply cos to odd indices in the array; 2i+1
   angle_rads[:, 1::2] = numpy.cos(angle_rads[:, 1::2])
   return jnp.array(angle_rads)

n_tokens_per_review = X_tokens.shape[1]
position_encoding = make_position_encoding(n_tokens_per_review, embed_dim)

plt.imshow(position_encoding[:, :]);
plt.axis('auto')

######################################################################
#  3. Make weight matrices for mapping word embedding to key, query, and value.
#  4. Make weight matrix for combining output of all heads

def make_weights(n_in, n_out):
    w_scale = 1 / jnp.sqrt(n_in)
    return jnp.array(numpy.random.uniform(-w_scale, w_scale, size=(n_in, n_out)))

n_heads = 8
n_in_per_head = embed_dim

W_keys = [make_weights(n_in_per_head, n_in_per_head) for h in range(n_heads)]
W_queries = [make_weights(n_in_per_head, n_in_per_head) for h in range(n_heads)]
W_values = [make_weights(n_in_per_head, n_in_per_head) for h in range(n_heads)]
W_combine = make_weights(n_heads * n_in_per_head, embed_dim)

print(f'W_key shapes {[w.shape for w in W_keys]}')
print(f'W_query shapes {[w.shape for w in W_queries]}')
print(f'W_value shapes {[w.shape for w in W_values]}')
print(f'W_combine shape {W_combine.shape}')

######################################################################
#  6. Make weights for dense net to apply to output of self-attention.

n_ff_units = 10
ff_W1 = make_weights(embed_dim, n_ff_units * embed_dim)
ff_W2 = make_weights(n_ff_units * embed_dim, embed_dim)

print(f'ff_W1 shape {ff_W1.shape}')
print(f'ff_W2 shape {ff_W2.shape}')

######################################################################
#  8. Make weights to linearly convert output of dense net to log probs for each class.

n_classes = 2
W_toprobs = make_weights(embed_dim, n_classes)
print(f'W_toprobs shape {W_toprobs.shape}')

######################################################################
#  Finally, start classification.
#  9. Convert list of tokens for each review into their respective embeddings.

embedding = jnp.take(embedder_W, X_tokens, axis=0)  # n_reviews x max_n_tokens x embed_dim
X = embedding + position_encoding
print(f'X.shape {X.shape}')

######################################################################
#  10. Pass each embedding through self-attention then dense net.

def forward(params, X_tokens):  # X is a  mini-batch    ,first, last):
    # Each row of X is an array of token indices, one for each word 
    embedder_W = params[0]
    X_embedding = jnp.take(embedder_W, X_tokens, axis=0)  # n_reviews x max_n_tokens x embed_dim
    X = X_embedding + position_encoding
    Y = forward_transform_block(params, forward_attention(params, X)[0], X)
    W_toprobs = params[-3]
    Y = Y @ W_toprobs
    # print(f'Output Y shape {Y.shape}')
    Y = Y.mean(axis=1)  # mean over all outputs
    return Y

params = [embedder_W, W_keys, W_queries, W_values, W_combine, W_toprobs, ff_W1, ff_W2]

forward(params, X_tokens)

labels.reshape(-1, 1) == numpy.unique(labels)

def make_indicator_vars(labels):
    return jnp.array((labels.reshape(-1, 1) == numpy.unique(labels)).astype(int))

make_indicator_vars(labels)

loss(params, X_tokens, T)

Xtrain, Ttrain, Xtest, Ttest = next(generate_stratified_partitions(X_tokens, T, 4, 
                                                                   validation=False, shuffle=True))
print(f'{Xtrain.shape=} {Ttrain.shape=} {Xtest.shape=} {Ttest.shape=}')

def frac_pos(T):
    return (T[:, 1] == 1).mean().item()
    
print(f'{frac_pos(Ttrain)=:.2f} {frac_pos(Ttest)=:.2f}')

batch_size = -1
learning_rate = 0.1 
n_steps = 500

losses, likelihoods = train(n_steps, batch_size, learning_rate)

plt.plot(likelihoods);

# To use our transformer, do a forward pass in batches, convert outputs to probabilities, then calculate a confusion matrix.

Y = []
n_samples = X_tokens.shape[0]
if batch_size > 0:
    first = 0
    for first in range(0, n_samples, batch_size):
        Y.append(forward(params, X_tokens[first:first + batch_size]))
    Y = jnp.vstack(Y)
else:
    Y = forward(params, X_tokens)

probs = softmax(Y, dim=1)
probs

pred_classes = jnp.argmax(Y, axis=1)
actual_classes = labels # jnp.argmax(labels, axis=1)

row0 = [jnp.mean(pred_classes[actual_classes == 0] == 0),
        jnp.mean(pred_classes[actual_classes == 0] == 1)]

row1 = [jnp.mean(pred_classes[actual_classes == 1] == 0),
        jnp.mean(pred_classes[actual_classes == 1] == 1)]
cm = jnp.array([row0, row1])

pandas.DataFrame(100*cm, 
                 columns=('Pred Neg', 'Pred Pos'),
                 index=('Actual Neg', 'Actual Pos'))

pandas.set_option('max_colwidth', None)

n = 10
df = pandas.DataFrame(numpy.hstack((actual_classes[:n].reshape(-1, 1),
                                    pred_classes[:n].reshape(-1, 1),
                                    headlines_orig[:n])),
                      columns=('Actual', 'Predicted', 'Headline'))
df

pandas.set_option('max_colwidth', None)

df = pandas.DataFrame(numpy.hstack((actual_classes[-n:].reshape(-1, 1),
                                    pred_classes[-n:].reshape(-1, 1),
                                    headlines_orig[-n:])),
                      columns=('Actual', 'Predicted', 'Headline'))
df

###### attention weights
headline_i = 1900

embedder_W = params[0]
X_embedding = jnp.take(embedder_W, X_tokens[headline_i:headline_i+1], axis=0)  # n_reviews x max_n_tokens x embed_dim
X = X_embedding + position_encoding  # [None, :, :]
attn, QKs = forward_attention(params, X)

words = headlines[headline_i]
n_words = words.index(' ') - 1
words = words[:n_words]
# n_words = len(headlines[headline_i])
plt.figure(figsize=(12, 12))
nplot = int(numpy.sqrt(n_heads)) + 1
ploti = 0
for h in range(n_heads):
    ploti += 1
    plt.subplot(nplot, nplot, ploti)
    plt.imshow(QKs[h][0, :n_words, :n_words])
    plt.colorbar()
plt.suptitle(' '.join(words))
plt.tight_layout()

## Draw attention weights as lines between words
plt.figure(figsize=(12, 12))

plt.suptitle(' '.join(words)) # headlines_orig[headline_i][0])
ploti = 0
for h in range(n_heads):
    ploti += 1
    plt.subplot(nplot, nplot, ploti)
    plt.xlim(0, 4)
    plt.ylim(0, n_words)
    plt.axis('off')
    for i, w in enumerate(words):
        plt.text(1, n_words - i, w, ha='right')
        plt.text(3, n_words - i, w, ha='left')

    QK = QKs[h][0].clone()[:n_words, :n_words]
    if True:
        mx = numpy.max(QK)
        # mn = numpy.min(QK)
        mask = QK < (0.7 * mx)
        QK = QK.at[mask].set(0.0)
        QK = QK / mx
    for i in range(n_words):
        for j in range(n_words):
            plt.plot([1, 3], [n_words - i, n_words - j], 'r', alpha=QK[i, j].item())
        
plt.tight_layout()